import re
from PyQt5.QtGui import QColor, QFont
from PyQt5.Qsci import QsciScintilla, QsciLexerCustom

class RstLexer(QsciLexerCustom):
    def __init__(self, parent=None):
        super().__init__(parent)
        
        # 定义所有样式及颜色
        self.styles = {
            0: "Default",     # 默认文本
            1: "Headers",     # 标题
            2: "Directives",  # 指令
            3: "Links",       # 链接
            4: "Lists",       # 列表
            5: "CodeBlocks",  # 代码块
            6: "InlineMath",  # 行内数学公式
            7: "MathBlock",   # 数学块
            8: "MathKeyword"  # 数学关键词
        }
        
        # 设置颜色
        self.setColor(QColor("#000000"), 0)  # 黑色
        self.setColor(QColor("#006699"), 1)  # 深蓝色
        self.setColor(QColor("#990099"), 2)  # 紫色
        self.setColor(QColor("#007700"), 3)  # 绿色
        self.setColor(QColor("#666600"), 4)  # 橄榄色
        self.setColor(QColor("#666666"), 5)  # 灰色
        self.setColor(QColor("#CC0000"), 6)  # 红色（行内数学）
        self.setColor(QColor("#990000"), 7)  # 暗红色（数学块）
        self.setColor(QColor("#009900"), 8)  # 绿色（数学关键词）

        # 设置字体
        font = QFont("Consolas", 12)
        self.setFont(font)
        
        # 预编译正则表达式（注意顺序优先级）
        self._compile_rules()
        
        # 初始化状态变量
        self.in_math_block = False  # 数学块状态
        self.in_code_block = False  # 代码块状态

    def _compile_rules(self):
        """编译所有正则规则（按优先级排序）"""
        self.rules = [
            # 数学公式（高优先级）
            (r':math:`.+?`', 6),               # 行内数学 :math:`...`
            (r'\\\((.*?)\\\)', 6),             # 行内 \(...\)
            
            # 特殊块处理
            (r'^\s*\.\. math::.*', 7, 'math_block_start'),  # math指令
            
            # 常规语法（较低优先级）
            (r'^[=~`:\-+*^#]{3,}$', 1),        # 标题下划线
            (r'^\.\. [a-zA-Z0-9_]+::', 2),     # 指令
            (r'`[^`]+`_?', 3),                 # 链接
            (r'^\s*[\*\-+] ', 4),              # 无序列表
            (r'^\s*\d+\. ', 4),               # 有序列表
            (r'^\s*::\s*$', 5),               # 代码块标记
            (r'^( {4,}|\t)', 5, 'code_block_start'),  # 代码块缩进
        ]
        
        # 数学关键词正则
        self.math_keywords = re.compile(
            r'\\[a-zA-Z]+\b|'
            r'\b(sin|cos|tan|log|exp|lim|sum|prod|int)\b|'
            r'\\[alpha|beta|gamma|delta]',
            re.UNICODE
        )
        
        # 预处理规则
        self.compiled_rules = []
        for rule in self.rules:
            pattern = re.compile(rule[0], re.UNICODE)
            style = rule[1]
            callback = rule[2] if len(rule)>2 else None
            self.compiled_rules.append( (pattern, style, callback) )

    def language(self):
        return "reStructuredText"

    def description(self, style):
        return self.styles.get(style, "")

    def styleText(self, start, end):
        # print(f"\nProcessing range: {start}-{end}")

        editor = self.editor()
        text = editor.text()[start:end]
        self.startStyling(start)
        
        # 分块处理逻辑
        pos = 0
        length = len(text)
        
        while pos < length:
            # 处理块状态
            if self.in_math_block:
                end_pos = text.find('$$', pos)
                if end_pos == -1:
                    self._apply_style(length - pos, 7)
                    pos = length
                else:
                    self._apply_style(end_pos - pos + 2, 7)
                    pos = end_pos + 2
                    self.in_math_block = False
                continue
                
            if self.in_code_block:
                # 代码块处理逻辑（略）
                pass
            
            # 匹配所有规则
            matched = False
            for pattern, style, callback in self.compiled_rules:
                match = pattern.match(text, pos)
                if match:
                    self._handle_match(match, style, callback)
                    pos = match.end()
                    matched = True
                    break
                    
            if not matched:
                # 默认样式
                self._apply_style(1, 0)
                pos += 1
        # print(f"Applied style {style} at position {pos} for {length} chars")

    def _handle_match(self, match, style, callback):
        """处理匹配到的规则"""
        text = match.group()
        
        # 执行回调处理特殊逻辑
        if callback == 'math_block_start':
            self.in_math_block = True
        elif callback == 'code_block_start':
            self.in_code_block = True
        
        # 应用基础样式
        self._apply_style(len(text), style)
        
        # 数学公式额外处理
        if style in (6, 7):
            self._highlight_math_keywords(text, style)

    def _apply_style(self, length, base_style):
        """应用基础样式"""
        if length > 0:
            self.setStyling(length, base_style)

    def _highlight_math_keywords(self, text, base_style):
        """在数学内容中高亮关键词"""
        last_pos = 0
        for match in self.math_keywords.finditer(text):
            start, end = match.start(), match.end()
            if start > last_pos:
                self.setStyling(start - last_pos, base_style)
            self.setStyling(end - start, 8)  # 关键词样式
            last_pos = end
        if last_pos < len(text):
            self.setStyling(len(text) - last_pos, base_style)